package org.example.memora.store;

import dev.langchain4j.data.message.*;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.example.memora.common.utils.ChatContext;
import org.example.memora.common.utils.ChatMessageUtils;
import org.example.memora.common.utils.UserHolder;
import org.example.memora.entity.ChatMessageDocument;
import org.example.memora.entity.ChatMessages;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.stream.Collectors;
import org.example.memora.DTO.UserRedisDTO;
import java.time.LocalDateTime;

@Component
public class MongoChatMemoryStore implements ChatMemoryStore {

    @Autowired
    private MongoTemplate mongoTemplate;

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        Query query = Query.query(Criteria.where("memoryId").is(String.valueOf(memoryId)))
                .with(Sort.by(Sort.Direction.ASC, "createTime"));

        List<ChatMessageDocument> docs = mongoTemplate.find(query, ChatMessageDocument.class);

        return docs.stream().map(doc -> {
            return switch (doc.getType()) {
                case "USER" -> new UserMessage(doc.getText());
                case "AI" -> new AiMessage(doc.getText());
                case "SYSTEM" -> new SystemMessage(doc.getText());
                //case "TOOL_EXECUTION_RESULT" -> new AiMessage(doc.getToolMessageText());
                default -> {
                    //System.out.println(">>> [getMessages] 未知类型: " + doc.getType());
                    yield null;
                }
            };
        }).filter(Objects::nonNull).collect(Collectors.toList());
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        //System.out.println(">>> updateMessages 被调用 memoryId=" + memoryId + "，消息条数=" + messages.size());
        String memoryIdStr = String.valueOf(memoryId);
        String userId = Optional.ofNullable(ChatContext.getCurrentUserId()).orElse("anonymous");

        mongoTemplate.remove(Query.query(Criteria.where("memoryId").is(memoryIdStr)), ChatMessageDocument.class);
        List<ChatMessageDocument> docs = new ArrayList<>();
        for (ChatMessage message : messages) {

            ChatMessageDocument doc = new ChatMessageDocument();
            //System.out.println(">>> [getMessages] 读取消息: type = " + doc.getType() + ", text = " + doc.getText());
            doc.setMemoryId(memoryIdStr);
            doc.setUserId(userId);
            doc.setCreateTime(LocalDateTime.now());

            if (message instanceof UserMessage user) {
                doc.setType("USER");
                String input = user.singleText();
                doc.setText(user.singleText());
                String keyword = "Answer using the following information:";
                int index = input.indexOf("Answer using the following information");
                // 保留关键词之前的部分
                if (index != -1) {
                    String res=input.substring(0, index).trim();
                    doc.setText(res);
                }else{
                    String res= input.trim();
                    doc.setText(res);
                }
                // 如果没找到关键词，返回原文本

            } else if (message instanceof AiMessage ai) {
                doc.setType("AI");
                doc.setText(ai.text());
            } else if (message instanceof SystemMessage sys) {
                doc.setType("SYSTEM");
                doc.setText(sys.text());
                //continue;
            }
            else {
                //System.out.println(">>> [updateMessages] 跳过未知类型: " + message);
                continue;
            }
            docs.add(doc);
        }

        mongoTemplate.insertAll(docs);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        Query query = Query.query(Criteria.where("memoryId").is(String.valueOf(memoryId)));
        mongoTemplate.remove(query, ChatMessageDocument.class);
    }
    public List<ChatMessageDocument> getUserAiMessagesPaged(String memoryId, String userId, int page, int pageSize) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId)
                .and("userId").is(userId)
                .and("type").in("USER", "AI")
                .and("text").not().in(null, "");

        Query query = new Query(criteria)
                .with(Sort.by(Sort.Direction.ASC, "createTime"))
                .skip((long) (page - 1) * pageSize)
                .limit(pageSize);

        return mongoTemplate.find(query, ChatMessageDocument.class);
    }

    public long countUserAiMessages(String memoryId, String userId) {
        Criteria criteria = Criteria.where("memoryId").is(memoryId)
                .and("userId").is(userId)
                .and("type").in("USER", "AI");

        return mongoTemplate.count(new Query(criteria), ChatMessageDocument.class);
    }

}

//@Component
//public class
//MongoChatMemoryStore implements ChatMemoryStore {
//
//    @Autowired
//    private MongoTemplate mongoTemplate;
//
//    @Override
//    public List<ChatMessage> getMessages(Object memoryId) {
//        Criteria criteria = Criteria.where("memoryId").is(memoryId);
//        Query query = new Query(criteria);
//        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);
//        if (chatMessages == null) {
//            return new LinkedList<>();
//        }
//        return ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
//    }
//
//    @Override
//    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
//
//        Criteria criteria = Criteria.where("memoryId").is(memoryId);
//        Query query=new Query(criteria);
//        Update update=new Update();
//        update.set("content", ChatMessageSerializer.messagesToJson(messages));
//
//        // 使用 upsert 方法，如果不存在则插入新文档
//        mongoTemplate.upsert(query, update, ChatMessages.class);
//
//    }
//
//    @Override
//    public void deleteMessages(Object memoryId) {
//
//        Criteria criteria = Criteria.where("memoryId").is(memoryId);
//        Query query=new Query(criteria);
//        mongoTemplate.remove(query, ChatMessages.class);
//
//    }
//}
